# coding=utf-8
import os
import requests
import tarfile
import zipfile

requests.packages.urllib3.disable_warnings()


def download_dataset(url, target_path):
    '''download and extract datasets'''
    if not os.path.exists(target_path):
        os.makedirs(target_path)
    download_file = os.path.join(target_path, url.split("/")[-1])
    if not os.path.exists(download_file):
        res = requests.get(url, stream=True, verify=False)
        if download_file.split(".")[-1] not in ['tgz', "zip", "tar", "gz"]:
            download_file = os.path.join(target_path, download_file)
        with open(download_file, "wb") as f:
            for chunck in res.iter_content(chunk_size=512):
                if chunck:
                    f.write(chunck)
        if download_file.endswith("zip"):
            z = zipfile.ZipFile(download_file, "r")
            z.extractall(path=target_path)
            z.close()
        if download_file.endswith(".tar.gz") or download_file.endswith(".tar") or download_file.endswith(".tgz"):
            t = tarfile.open(download_file)
            names = t.getnames()
            for name in names:
                t.extract(name, target_path)
            t.close()
        print("The {} file is downloaded and saved in the path {} after processing".format(os.path.basename(url),
                                                                                           target_path))


download_dataset("https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cifar-10-binary.tar.gz",
                 "../../datasets")
download_dataset("https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/MNIST_Data.zip",
                 "../../datasets")


import mindspore.dataset as ds

DATA_DIR = "../../datasets/cifar-10-batches-bin"
sampler = ds.SequentialSampler(num_samples=5)
dataset = ds.Cifar10Dataset(DATA_DIR, sampler=sampler)
for data in dataset.create_dict_iterator():
    print("Image shape: {}".format(data['image'].shape), ", Label: {}".format(data['label']))

# 自定义数据集
# 对于目前MindSpore不支持直接加载的数据集，可以构造自定义数据集类，然后通过GeneratorDataset接口实现自定义方式的数据加载。
import numpy as np

np.random.seed(58)


class DatasetGenerator:
    def __init__(self):
        self.data = np.random.sample((5, 2))
        self.label = np.random.sample((5, 1))

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return len(self.data)


# 定义数据集类之后，就可以通过GeneratorDataset接口按照用户定义的方式加载并访问数据集样本。
dataset_generator = DatasetGenerator()
dataset = ds.GeneratorDataset(dataset_generator, ["data", "label"], shuffle=False)

for data in dataset.create_dict_iterator():
    print('{}'.format(data["data"]), '{}'.format(data["label"]))
ds.config.set_seed(58)

# 数据处理
# 随机打乱数据顺序
dataset = dataset.shuffle(buffer_size=10)
# 对数据集进行分批
dataset = dataset.batch(batch_size=2)

for data in dataset.create_dict_iterator():
    print("data: {}".format(data["data"]))
    print("label: {}".format(data["label"]))


# 数据增强
import matplotlib
import matplotlib.pyplot as plt
from mindspore.dataset.vision import Inter
import mindspore.dataset.vision.c_transforms as c_vision


DATA_DIR = '../../datasets/MNIST_Data/train'

mnist_dataset = ds.MnistDataset(DATA_DIR, num_samples=6, shuffle=False)

# 查看数据原图
mnist_it = mnist_dataset.create_dict_iterator()
data = next(mnist_it)
matplotlib.use('TkAgg')
plt.figure(figsize=(3, 3))
plt.imshow(data['image'].asnumpy().squeeze(), cmap=plt.cm.gray)
plt.title(data['label'].asnumpy(), fontsize=20)
plt.show()
# 定义数据增强算子，对数据集进行Resize和RandomCrop操作，然后通过map映射将其插入数据处理管道。
resize_op = c_vision.Resize(size=(200, 200), interpolation=Inter.LINEAR)
crop_op = c_vision.RandomCrop(150)
transforms_list = [resize_op, crop_op]
mnist_dataset = mnist_dataset.map(operations=transforms_list, input_columns=["image"])

# 查看数据增强效果。
mnist_dataset = mnist_dataset.create_dict_iterator()
data = next(mnist_dataset)
plt.figure(figsize=(3, 3))
plt.imshow(data['image'].asnumpy().squeeze(), cmap=plt.cm.gray)
plt.title(data['label'].asnumpy(), fontsize=20)
plt.show()
